package miao.byusi.q0j;

import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;

import java.io.*;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.util.*;
import java.util.zip.ZipEntry;
import java.util.zip.ZipInputStream;
import java.util.zip.ZipOutputStream;

public class App {
    private static final Logger logger = LogManager.getLogger(App.class);

    public static void main(String[] args) {
        if (args.length < 2) {
            logger.error("参数不足，用法: java -jar q0jm.jar [-j <文件>] 或 [-m <ZIP文件>]");
            return;
        }

        String mode = args[0];
        String filePath = args[1];
        try {
            switch (mode) {
                case "-j":
                    encrypt(filePath);
                    break;
                case "-m":
                    decrypt(filePath);
                    break;
                default:
                    logger.error("无效的模式，请使用 -j（加密）或 -m（解密）");
            }
        } catch (Exception e) {
            logger.error("操作失败: {}", e.getMessage(), e);
        }
    }

    private static void encrypt(String sourceFilePath) throws IOException {
        logger.info("开始加密文件: {}", sourceFilePath);
        File sourceFile = new File(sourceFilePath);
        if (!sourceFile.exists()) {
            throw new FileNotFoundException("文件未找到: " + sourceFilePath);
        }

        // 读取文件并处理零字节
        logger.debug("读取源文件内容...");
        byte[] originalData = Files.readAllBytes(sourceFile.toPath());
        List<Integer> zeroPositions = new ArrayList<>();
        ByteArrayOutputStream nonZeroOut = new ByteArrayOutputStream();

        logger.info("正在扫描并移除零字节...");
        for (int i = 0; i < originalData.length; i++) {
            byte b = originalData[i];
            if (b == 0) {
                zeroPositions.add(i);
            } else {
                nonZeroOut.write(b);
            }
        }

        // 生成 q0jm.qz
        logger.debug("生成 q0jm.qz 文件...");
        byte[] nonZeroBytes = nonZeroOut.toByteArray();
        ByteArrayOutputStream qzBaos = new ByteArrayOutputStream();
        qzBaos.write("q0jm".getBytes(StandardCharsets.US_ASCII)); // 写入文件头
        qzBaos.write(nonZeroBytes);
        byte[] qzData = qzBaos.toByteArray();

        // 生成 zre.json
        logger.debug("生成 zre.json 元数据...");
        Map<String, Object> zreMap = new HashMap<>();
        zreMap.put("originalFileName", sourceFile.getName());
        zreMap.put("zeroPositions", zeroPositions);
        ByteArrayOutputStream jsonBaos = new ByteArrayOutputStream();
        new ObjectMapper().writeValue(jsonBaos, zreMap);
        byte[] jsonData = jsonBaos.toByteArray();

        // 打包为 ZIP
        String zipFileName = sourceFile.getName() + ".zip";
        logger.info("正在生成ZIP文件: {}", zipFileName);
        try (ZipOutputStream zos = new ZipOutputStream(new FileOutputStream(zipFileName))) {
            zos.putNextEntry(new ZipEntry("q0jm.qz"));
            zos.write(qzData);
            zos.closeEntry();

            zos.putNextEntry(new ZipEntry("zre.json"));
            zos.write(jsonData);
            zos.closeEntry();
        }

        logger.info("加密成功! ZIP文件已保存为: {}", zipFileName);
        logger.info("原始文件大小: {} 字节, 移除零字节数: {}, 有效数据大小: {} 字节",
                originalData.length, zeroPositions.size(), nonZeroBytes.length);
    }

    private static void decrypt(String zipFilePath) throws IOException {
        logger.info("开始解密ZIP文件: {}", zipFilePath);
        File zipFile = new File(zipFilePath);
        if (!zipFile.exists()) {
            throw new FileNotFoundException("ZIP文件未找到: " + zipFilePath);
        }

        // 从ZIP提取数据
        logger.debug("解析ZIP文件内容...");
        byte[] qzData = null;
        byte[] jsonData = null;

        try (ZipInputStream zis = new ZipInputStream(new FileInputStream(zipFile))) {
            ZipEntry entry;
            while ((entry = zis.getNextEntry()) != null) {
                ByteArrayOutputStream baos = new ByteArrayOutputStream();
                byte[] buffer = new byte[1024];
                int len;
                while ((len = zis.read(buffer)) > 0) {
                    baos.write(buffer, 0, len);
                }
                if (entry.getName().equals("q0jm.qz")) {
                    qzData = baos.toByteArray();
                } else if (entry.getName().equals("zre.json")) {
                    jsonData = baos.toByteArray();
                }
                zis.closeEntry();
            }
        }

        // 校验数据完整性
        if (qzData == null || jsonData == null) {
            throw new IOException("ZIP文件缺少必要内容: q0jm.qz 或 zre.json");
        }

        // 解析 q0jm.qz
        logger.debug("解析 q0jm.qz 数据...");
        if (qzData.length < 4 || !"q0jm".equals(new String(qzData, 0, 4, StandardCharsets.US_ASCII))) {
            throw new IOException("无效的 q0jm.qz 文件头");
        }
        byte[] nonZeroBytes = Arrays.copyOfRange(qzData, 4, qzData.length);

        // 解析 zre.json
        logger.debug("解析 zre.json 元数据...");
        ObjectMapper mapper = new ObjectMapper();
        Map<String, Object> zreMap = mapper.readValue(jsonData, new TypeReference<Map<String, Object>>() {});
        String originalFileName = (String) zreMap.get("originalFileName");
        List<?> rawZeroPositions = (List<?>) zreMap.get("zeroPositions");

        // 处理位置类型
        List<Integer> zeroPositions = new ArrayList<>();
        for (Object posObj : rawZeroPositions) {
            if (posObj instanceof Number) {
                zeroPositions.add(((Number) posObj).intValue());
            } else {
                throw new IOException("无效的位置类型: " + posObj.getClass());
            }
        }

        // 恢复原始数据
        logger.info("正在恢复文件: {}", originalFileName);
        int originalLength = nonZeroBytes.length + zeroPositions.size();
        byte[] originalData = new byte[originalLength];
        Set<Integer> zeroSet = new HashSet<>(zeroPositions);

        int nonZeroIndex = 0;
        for (int i = 0; i < originalLength; i++) {
            if (zeroSet.contains(i)) {
                originalData[i] = 0;
            } else {
                if (nonZeroIndex >= nonZeroBytes.length) {
                    throw new IOException("数据错误: 非零字节不足");
                }
                originalData[i] = nonZeroBytes[nonZeroIndex++];
            }
        }

        // 写入文件
        Files.write(new File(originalFileName).toPath(), originalData);
        logger.info("解密成功! 文件已保存为: {}", originalFileName);
        logger.info("恢复后的文件大小: {} 字节, 填充零字节数: {}", originalLength, zeroPositions.size());
    }
}